from crawl4ai.extraction_strategy import JsonCssExtractionStrategy
from enum import Enum


class ExtractionStrategyVersion(str, Enum):
    V1 = "v1"
    V2 = "v2"


def get_ExtractionStrategy(
    ver: ExtractionStrategyVersion = ExtractionStrategyVersion.V2,
) -> JsonCssExtractionStrategy:
    print(f"get_ExtractionStrategy use ver:{ver.value}")
    if ver == ExtractionStrategyVersion.V1:
        schema = {
            "name": "SongLyrics_Page",
            "baseSelector": "div.left",
            "fields": [
                {"name": "title", "selector": "#sonsyuanwen h1", "type": "text"},
                {
                    "name": "author",
                    "selector": "#sonsyuanwen p.source > a:nth-child(1)",
                    "type": "text",
                },
                {
                    "name": "dynasty",
                    "selector": "#sonsyuanwen p.source > a:nth-child(2)",
                    "type": "text",
                },
                {
                    "name": "content",
                    "selector": "#sonsyuanwen div.contson",
                    "type": "text",
                },
                {
                    "name": "translation",
                    "selector": "div[id^=fanyiquan] > div.contyishang > p:nth-child(2)",
                    "type": "text",
                },
                {
                    "name": "comment",
                    "selector": "div[id^=fanyiquan] > div.contyishang > p:nth-child(3)",
                    "type": "text",
                },
                {
                    "name": "appreciation",
                    "selector": "div[id^=shangxiquan] > div.contyishang > p",
                    "type": "list",
                    "fields": [{"name": "value", "type": "text"}],
                },
                {
                    "name": "creative_background",
                    "selector": "div.sons:nth-child(7) > div.contyishang > p",
                    "type": "text",
                },
                {
                    "name": "brief_analysis",
                    "selector": "div.sons:nth-child(8) > div.contyishang > p",
                    "type": "list",
                    "fields": [{"name": "value", "type": "text"}],
                },
            ],
        }
    elif ver == ExtractionStrategyVersion.V2:
        schema = {
            "name": "Detail_Page",
            "baseSelector": "div.left",
            "fields": [
                {"name": "标题", "selector": "#sonsyuanwen h1", "type": "text"},
                {
                    "name": "作者",
                    "selector": "#sonsyuanwen p.source > a:nth-child(1)",
                    "type": "text",
                },
                {
                    "name": "朝代",
                    "selector": "#sonsyuanwen p.source > a:nth-child(2)",
                    "type": "text",
                },
                {
                    "name": "内容",
                    "selector": "#sonsyuanwen div.contson",
                    "type": "text",
                },
                {
                    "name": "译文及注释",
                    "selector": "div.sons:not([style*='none']):has(div.contyishang a[href*='译文及注释']) > div.contyishang > p",
                    "type": "list",
                    "fields": [{"name": "value", "type": "text"}],
                },
                {
                    "name": "赏析",
                    "selector": "div.sons:not([style*='none']):has(div.contyishang a[href*='赏析'], div.contyishang a[href*='鉴赏']) > div.contyishang > p",
                    "type": "list",
                    "fields": [{"name": "value", "type": "text"}],
                },
                {
                    "name": "创作背景",
                    "selector": "div.sons:not([style*='none']):has(div.contyishang a[href*='创作背景']) > div.contyishang > p",
                    "type": "list",
                    "fields": [{"name": "value", "type": "text"}],
                },
                {
                    "name": "简析",
                    "selector": "div.sons:not([style*='none']):has(div.contyishang a[href*='简析']) > div.contyishang > p",
                    "type": "list",
                    "fields": [{"name": "value", "type": "text"}],
                },
            ],
        }
    return JsonCssExtractionStrategy(schema)


## div.sons:not([style*='none']):has(div.contyishang a[href*='译文及注释']) > div.contyishang > p
## div.sons:not([style*='none']):has(div.contyishang a[href*='赏析'], div.contyishang a[href*='鉴赏']) > div.contyishang > p
## div.sons:not([style*='none']):has(div.contyishang a[href*='创作背景']) > div.contyishang > p
## div.sons:not([style*='none']):has(div.contyishang a[href*='简析']) > div.contyishang > p
